import io
import time
import joblib
import pandas as pd
from sklearn.pipeline import Pipeline

_MODEL_PATH = "filing_rate_2/model.pkl"
_pipe: Pipeline = None


def load_model():
    """加载模型"""
    global _pipe
    try:
        _pipe = joblib.load(_MODEL_PATH)
        print("模型加载成功")
    except Exception as e:
        print(f"模型加载失败: {e}")
        raise


def predict(raw: list[dict]) -> list[float]:
    """预测立案概率"""
    if _pipe is None:
        load_model()

    df = pd.DataFrame(raw)
    return _pipe.predict_proba(df)[:, 1].tolist()


def update(raw: list[dict], y: list[int]):
    """增量训练"""
    if _pipe is None:
        load_model()

    df = pd.DataFrame(raw)
    # 只更新分类器，预处理复用
    X = _pipe.named_steps["prep"].transform(df)
    _pipe.named_steps["clf"].partial_fit(X, y)

    # 保存更新后的模型
    joblib.dump(_pipe, _MODEL_PATH)
    print(f"增量训练完成，更新了 {len(y)} 个样本")


def update_csv(csv: bytes):
    """增量训练，从 CSV 数据中读取"""
    df = pd.read_csv(io.BytesIO(csv))
    raw = df.to_dict(orient="records")
    y = df["result"].tolist()
    update(raw, y)


if __name__ == "__main__":
    raw_1 = [
        {
            "mid": 1504,
            "amount": 125519.07,
            "btype": 5,
            "sx_state": 0,
            "zx_amount": 0,
            "area": 540329,
            "field_standard": 1,
            "court": 2983,
            "lbh_label": 0,
            "med_time": 0,
            "y_state": 1,
            "call_yjt_count": 4,
            "call_wjt_count": 4,
            "call_yjt_bg": 0,
            "call_wjt_bg": 1,
            "zx_to_tj_days": 0.0,
            "add_to_tj_days": 6.342523148148148,
            "tj_to_lbh_days": 0.8879050925925925,
        },
    ]
    raw_2 = [
        {
            "mid": 1504,
            "amount": 23165.3,
            "btype": 5,
            "sx_state": 0,
            "zx_amount": 0,
            "area": 140107,
            "field_standard": 1,
            "court": 1978,
            "lbh_label": 25,
            "med_time": 0,
            "y_state": 0,
            "call_yjt_count": 0,
            "call_wjt_count": 0,
            "call_yjt_bg": 0,
            "call_wjt_bg": 0,
            "zx_to_tj_days": 0.0,
            "add_to_tj_days": 18.123958333333334,
            "tj_to_lbh_days": 0.8879050925925925,
        },
    ]
    # with open("filing_rate/data.csv", "rb") as f:
    #     csv = f.read()
    # update_csv(csv)
    start = time.time()
    result = predict(raw_1)
    end = time.time()
    print(result[0])
    print(f"预测耗时: {end - start} 秒")
